# 导入onnx库，用于处理ONNX模型
import onnx
# 导入torch库，用于深度学习模型的构建和操作
import torch
# 从onnxsim库中导入simplify函数，用于简化ONNX模型
from onnxsim import simplify

# 注意：ckpt_file 和 args 在原代码中未定义，需要确保在使用前已经正确赋值
# 加载模型的状态字典，将模型加载到CPU上
ckpt = torch.load(ckpt_file, map_location="cpu")
# 将加载的状态字典应用到模型上
model.load_state_dict(ckpt)

# 准备一个虚拟输入，用于导出模型
# 虚拟输入的形状为 (batch_size, 3, exp.test_size[0], exp.test_size[1])
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])

# 导出模型为ONNX格式
torch.onnx._export(
    # 要导出的模型
    model,
    # 模型的输入，这里使用虚拟输入
    dummy_input,
    # 导出的ONNX模型文件名
    "yolox.onnx",
    # 输入张量的名称
    input_names=["input"],
    # 输出张量的名称
    output_names=["output"],
    # 定义动态轴，这里将输入和输出的第一个维度设置为动态的批次大小
    dynamic_axes={"input": {0: 'batch'},
                  "output": {0: 'batch'}})

# 使用onnx-simplifier库来简化导出的ONNX模型
# 加载导出的ONNX模型
onnx_model = onnx.load(args.output_name)
# 简化加载的ONNX模型，返回简化后的模型和验证结果
model_simp, check = simplify(onnx_model)
# 断言简化后的模型通过验证，如果验证失败，程序将抛出异常
assert check, "Simplified ONNX model could not be validated"
# 将简化后的模型保存到指定的输出文件中
onnx.save(model_simp, args.output_name)